'''
如果date格式缺少20，运行rename date脚本
'''

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import seaborn as sns
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,split_str_id,tydat_CMA,load_land_polygons,detect_landfall
import matplotlib.ticker as ticker
from tqdm import tqdm
from typlot.config.global_config import *

all_ini_time_mode = ['00_12','00','12']
names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
tynames,tyids = split_str_id(names)
all_obs_baseline= ['RI','land']  # ‘land’ 'RI'
draw_opt = False

for ini_time_mode in all_ini_time_mode:
    for obs_baseline in all_obs_baseline:             
        for tyname,tyid  in zip(tynames,tyids):
            '''settings'''
            # tyname,tyid = 'dusurui','16'
            dir_dsr=os.path.join(global_ensdir,f"{tyname}_{tyid}")
            name_date=sorted(os.listdir(dir_dsr))
            name_date = [ i for i in name_date if i[-2:] in ini_time_mode ]
            dir_date= [os.path.join(dir_dsr,f) for f in name_date if f[-2:] in ini_time_mode]
            RIstd = 7  # std = 15   
            member_num = 52
            draw_obs_opt = True
            obs_path=os.path.join(global_obsdir,f"{tyname}_CMAobs.txt")
            pic_savepath=os.path.join(global_picdir,f"{obs_baseline}_heatmap",ini_time_mode,"start_day_RIsum")
            os.makedirs(pic_savepath,exist_ok=True)
            
            
            ###  记录下所有数据  ###
            num_array = np.zeros((len(name_date),member_num),dtype=int)
            #用于记录所有数据；所有起报时间内的所有成员以及对应成员ID
            num=[]
            for date in dir_date:
                # 得到某一起报时刻的所有集合成员数据
                name_all = [f for f in os.listdir(date) if f.startswith("TRACK")]
                sorted_names = sorted(name_all, key=lambda x: int(x.split('TRACK_ID_')[-1]))
                path_all = [os.path.join(date, f) for f in sorted_names ]
                num2=[]
                #通过列表记录所有成员，借助类来封装处理方式
                for path in path_all:
                    t = tydat(path,RIstd)
                    nn = int(path.split("TRACK_ID_")[-1] )
                    # 记录单个成员数据以及成员序号
                    num2.append([t,nn]) 
                num.append(num2)
            
            
            
            
            ###  创建绘图数组   ###
            a=[];b=[]
            for s in range(0,len(name_date)):
                start_time=datetime.strptime(name_date[s],"%Y%m%d%H")
                for i in range(0,len(num[s])):
                    a.append( max(num[s][i][0].time ))
                    b.append( max(num[s][i][0].time) - start_time )
            max_enddate = max(a)
            max_dayrange = max(b)
            draw_data = np.zeros( (len(name_date),max_dayrange.days+1) ,dtype=int)
            
            
            
            ''' fill members draw array '''
            tt,dd = draw_data.shape
            #起报
            for t in tqdm(range(tt),desc=f'processing {tyname} {ini_time_mode} {obs_baseline}' ):
                #预报时效
                for d in range(dd):
                    start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                    cal_time = start_time + d*timedelta(days=1)
                    calend_time = cal_time + timedelta(hours=18)
                    #成员
                    n=0     
                    for i in range(len(num[t])): #对成员数循环
                        data = num[t][i][0]
                        RI = data.num_rapidgrow()
                        dayRI = RI[ (data.time>=cal_time) & (data.time<=calend_time) ]
                        if np.sum(dayRI>=1) :
                            n+=1
                    draw_data[t][d]=n
                        
                
            
            if draw_obs_opt == True : 
               ''' init obsCMA; RI points reord '''
               obs = tydat_CMA(obs_path) 
               draw_data_obs = np.zeros( (len(name_date),max_dayrange.days+1),dtype=int )
               tt,dd = draw_data_obs.shape
               if obs_baseline == 'RI':
                   record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)
                   for t in range(tt):
                       #预报时效
                       for d in range(dd):
                           try:
                               start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                               cal_time = start_time + d*timedelta(days=1)
                               calend_time = cal_time + timedelta(days=1)
                               # 对齐
                               trange = (obs.time>=cal_time) & (obs.time<calend_time)
                               draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 
                           except Exception as e:
                               print('line121',e)           
               ''' iter to all start time '''
               if obs_baseline == 'land':
                    # get  land  time
                    land_polys = load_land_polygons(os.path.join(global_shpdir,'China','bou1_4p.shp'), os.path.join(global_shpdir,'China','bou1_4p.dbf'))
                    flags,id1 = detect_landfall(obs.lat, obs.lon, land_polys)
                    land_time = obs.time[id1]
                    # identify  position
                    for t in range(tt):
                        for d in range(dd):
                            try:
                                start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                                cal_time = start_time + d*timedelta(days=1)
                                calend_time = cal_time + timedelta(days=1)
                                trange = (obs.time>=cal_time) & (obs.time<calend_time)
                                # 填充
                                draw_data_obs[t][d] = 1 if land_time in obs.time[trange] else 0
                            except Exception as e:
                                print(e)
               

            
            fig, ax = plt.subplots(figsize=(15, 9)) 
            
            ''' heatmap '''
            ax = sns.heatmap(draw_data[:,:15],annot=True,
                             annot_kws={'size': 6},fmt="d", cmap="YlGnBu", 
                             cbar_kws={'label': 'number of rapid intensification'}, 
                             linewidths=0.5, ax=ax)  #,vmin=0,vmax =30
            
            ax.set_xlabel("Forecast Lead Time", fontsize=12)
            ax.set_ylabel("Forecast Initialization Time", fontsize=12)
            #  手动设置xticks数量,snsheatmap default上限是19个
            ax.set_yticks(np.arange(draw_data.shape[0])) 
            ax.set_yticklabels(name_date, rotation=0) 
            plt.xticks(rotation=0)
            xlabels=[str(i)+"d"  for i in range(1,16)]
            ax.set_xticklabels(xlabels)
            ax.set_title(tyname+f"-RI{RIstd} RIpoints number in all members   forecast_ini_mode:{ini_time_mode}" )
            cbar = ax.collections[0].colorbar
            cbar.locator = ticker.MaxNLocator(integer=True)
            cbar.update_ticks()
            
            
            ''' plot '''
            points = np.where(draw_data_obs[:,:15].T == 1)
            x_coords = points[0]+0.5
            y_coords = points[1]+0.15
            ax.scatter(x_coords, y_coords, color='red', marker='o',s=10)
        
            ax.invert_yaxis()
            if draw_opt == False:
                plt.savefig(os.path.join(pic_savepath,f"{tyname}.png"),dpi=900)
                plt.close()
            else:
                plt.show()
                
                
            plt.clf()
            plt.close('all')

